package main

import (
	"database/sql"
	"fmt"
	"log"
	"sync"

	_ "github.com/go-sql-driver/mysql"
)

// MySQLServer 处理MySQL请求的服务器实现
type MySQLServer struct {
	connections map[string]*MySQLHandler
	txMap       map[string]*sql.Tx
	stmtMap     map[string]*sql.Stmt
	txMutex     sync.RWMutex
	stmtMutex   sync.RWMutex
}

// NewMySQLServer 创建新的MySQL服务器
func NewMySQLServer(connections map[string]string) (*MySQLServer, error) {
	server := &MySQLServer{
		connections: make(map[string]*MySQLHandler),
		txMap:       make(map[string]*sql.Tx),
		stmtMap:     make(map[string]*sql.Stmt),
	}

	// 初始化所有MySQL连接
	for name, dsn := range connections {
		handler, err := NewMySQLHandler(dsn)
		if err != nil {
			return nil, fmt.Errorf("无法连接到MySQL '%s': %v", name, err)
		}
		server.connections[name] = handler
		log.Printf("MySQL连接 '%s' 已初始化", name)
	}

	return server, nil
}

// HandleRequest 处理MySQL请求
func (s *MySQLServer) HandleRequest(req SQLRequest) SQLResponse {
	// 获取MySQL处理器
	handler, exists := s.connections[req.Connection]
	if !exists {
		return SQLResponse{Error: "MySQL连接不存在: " + req.Connection}
	}

	// 根据操作类型执行不同的处理
	switch req.Operation {
	case "query":
		return s.handleQuery(handler, req)
	case "exec":
		return s.handleExec(handler, req)
	case "prepare":
		return s.handlePrepare(handler, req)
	case "begin":
		return s.handleBegin(handler, req)
	case "commit":
		return s.handleCommit(req)
	case "rollback":
		return s.handleRollback(req)
	default:
		return SQLResponse{Error: "未知操作类型: " + req.Operation}
	}
}

// 处理MySQL查询
func (s *MySQLServer) handleQuery(handler *MySQLHandler, req SQLRequest) SQLResponse {
	var rows *sql.Rows
	var err error

	// 检查是否在事务中
	if req.RequestID != "" {
		s.txMutex.RLock()
		tx, exists := s.txMap[req.RequestID]
		s.txMutex.RUnlock()

		if exists {
			rows, err = handler.QueryTx(tx, req.SQL, req.Args)
		} else {
			return SQLResponse{Error: "事务ID不存在"}
		}
	} else {
		rows, err = handler.Query(req.SQL, req.Args)
	}

	if err != nil {
		return SQLResponse{Error: err.Error()}
	}
	defer rows.Close()

	// 处理查询结果
	resultRows, columns, err := processSQLRows(rows)
	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	return SQLResponse{
		Columns: columns,
		Rows:    resultRows,
	}
}

// 处理MySQL执行
func (s *MySQLServer) handleExec(handler *MySQLHandler, req SQLRequest) SQLResponse {
	var result sql.Result
	var err error

	// 检查是否在事务中
	if req.RequestID != "" {
		s.txMutex.RLock()
		tx, exists := s.txMap[req.RequestID]
		s.txMutex.RUnlock()

		if exists {
			result, err = handler.ExecTx(tx, req.SQL, req.Args)
		} else {
			return SQLResponse{Error: "事务ID不存在"}
		}
	} else {
		result, err = handler.Exec(req.SQL, req.Args)
	}

	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	lastInsertID, _ := result.LastInsertId()
	rowsAffected, _ := result.RowsAffected()

	return SQLResponse{
		Rows:    [][]interface{}{{lastInsertID, rowsAffected}},
		Columns: []string{"last_insert_id", "rows_affected"},
	}
}

// 处理MySQL预处理语句
func (s *MySQLServer) handlePrepare(handler *MySQLHandler, req SQLRequest) SQLResponse {
	var stmt *sql.Stmt
	var err error

	// 检查是否在事务中
	if req.RequestID != "" {
		s.txMutex.RLock()
		tx, exists := s.txMap[req.RequestID]
		s.txMutex.RUnlock()

		if exists {
			stmt, err = handler.PrepareTx(tx, req.SQL)
		} else {
			return SQLResponse{Error: "事务ID不存在"}
		}
	} else {
		stmt, err = handler.Prepare(req.SQL)
	}

	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	stmtID := fmt.Sprintf("%s:%p", req.RequestID, stmt)
	s.stmtMutex.Lock()
	s.stmtMap[stmtID] = stmt
	s.stmtMutex.Unlock()

	return SQLResponse{
		Rows:    [][]interface{}{{stmtID}},
		Columns: []string{"stmt_id"},
	}
}

// 处理MySQL开始事务
func (s *MySQLServer) handleBegin(handler *MySQLHandler, req SQLRequest) SQLResponse {
	if req.RequestID == "" {
		return SQLResponse{Error: "事务需要RequestID"}
	}

	s.txMutex.Lock()
	defer s.txMutex.Unlock()

	if _, exists := s.txMap[req.RequestID]; exists {
		return SQLResponse{Error: "事务ID已存在"}
	}

	tx, err := handler.Begin()
	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	s.txMap[req.RequestID] = tx
	return SQLResponse{Rows: [][]interface{}{{"事务已开始"}}}
}

// 处理MySQL提交事务
func (s *MySQLServer) handleCommit(req SQLRequest) SQLResponse {
	if req.RequestID == "" {
		return SQLResponse{Error: "需要提供事务ID"}
	}

	s.txMutex.Lock()
	defer s.txMutex.Unlock()

	tx, exists := s.txMap[req.RequestID]
	if !exists {
		return SQLResponse{Error: "事务ID不存在"}
	}

	err := tx.Commit()
	delete(s.txMap, req.RequestID)

	if err != nil {
		return SQLResponse{Error: err.Error()}
	}
	return SQLResponse{Rows: [][]interface{}{{"事务已提交"}}}
}

// 处理MySQL回滚事务
func (s *MySQLServer) handleRollback(req SQLRequest) SQLResponse {
	if req.RequestID == "" {
		return SQLResponse{Error: "需要提供事务ID"}
	}

	s.txMutex.Lock()
	defer s.txMutex.Unlock()

	tx, exists := s.txMap[req.RequestID]
	if !exists {
		return SQLResponse{Error: "事务ID不存在"}
	}

	err := tx.Rollback()
	delete(s.txMap, req.RequestID)

	if err != nil {
		return SQLResponse{Error: err.Error()}
	}
	return SQLResponse{Rows: [][]interface{}{{"事务已回滚"}}}
}

// MySQL数据库处理器
type MySQLHandler struct {
	db *sql.DB
}

// 创建MySQL处理器
func NewMySQLHandler(dsn string) (*MySQLHandler, error) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}

	if err := db.Ping(); err != nil {
		return nil, err
	}

	return &MySQLHandler{
		db: db,
	}, nil
}

// 关闭连接
func (h *MySQLHandler) Close() error {
	if h.db != nil {
		return h.db.Close()
	}
	return nil
}

// 执行查询
func (h *MySQLHandler) Query(sql string, args []interface{}) (*sql.Rows, error) {
	return h.db.Query(sql, args...)
}

// 执行命令
func (h *MySQLHandler) Exec(sql string, args []interface{}) (sql.Result, error) {
	return h.db.Exec(sql, args...)
}

// 预处理语句
func (h *MySQLHandler) Prepare(sql string) (*sql.Stmt, error) {
	return h.db.Prepare(sql)
}

// 开始事务
func (h *MySQLHandler) Begin() (*sql.Tx, error) {
	return h.db.Begin()
}

// 在事务中查询
func (h *MySQLHandler) QueryTx(tx *sql.Tx, sql string, args []interface{}) (*sql.Rows, error) {
	return tx.Query(sql, args...)
}

// 在事务中执行
func (h *MySQLHandler) ExecTx(tx *sql.Tx, sql string, args []interface{}) (sql.Result, error) {
	return tx.Exec(sql, args...)
}

// 在事务中预处理
func (h *MySQLHandler) PrepareTx(tx *sql.Tx, sql string) (*sql.Stmt, error) {
	return tx.Prepare(sql)
}

// 处理MySQL查询结果
func processSQLRows(rows *sql.Rows) ([][]interface{}, []string, error) {
	columns, err := rows.Columns()
	if err != nil {
		return nil, nil, err
	}

	// 准备存储行数据的切片
	values := make([]interface{}, len(columns))
	scanArgs := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	// 读取所有行
	var allRows [][]interface{}
	for rows.Next() {
		err = rows.Scan(scanArgs...)
		if err != nil {
			return nil, nil, err
		}

		// 复制当前行的数据
		row := make([]interface{}, len(columns))
		for i, v := range values {
			row[i] = v
		}
		allRows = append(allRows, row)
	}

	if err = rows.Err(); err != nil {
		return nil, nil, err
	}

	return allRows, columns, nil
}
